﻿#include "RotateTemplateMatch.h"
#include "ImageRotate.h"
#include "TemplateMatch.h"
// #include <QDebug>
#include <iostream>
#include "opencv2/highgui.hpp"
#include "omp.h"

const int ANGLE_COUNT_MAX = 4096;

namespace Qly {

bool drawRotatedRect(cv::Mat &image, cv::RotatedRect &rrt, cv::Scalar color, int thickness)
{
    if(image.empty()) return false;
    //绘制旋转矩形
    cv::Point2f vertices[4];
    rrt.angle = -rrt.angle;
    rrt.points(vertices);
    rrt.angle = -rrt.angle;
    for (int i = 0; i < 4; i++)
    {
        cv::line(image, vertices[i], vertices[(i + 1) % 4], color, thickness);
    }
    cv::circle(image, vertices[1], thickness * 3, color, thickness);
    return true;
}


bool RotateTemplateMatch::better(double oldGrade, double newGrade, int method)
{
    switch (method)
    {
    case cv::TM_SQDIFF:
    case cv::TM_SQDIFF_NORMED:
        if(newGrade < oldGrade)
        {
            return true;
        }
         break;
    default:
        if(newGrade > oldGrade)
        {
            return true;
        }
    }
    return false;
}


bool RotateTemplateMatch::setModelByAngleCount(const cv::Mat &model, double angleMin, double angleMax, size_t angleCount)
{
    cv::Mat mask = cv::Mat(model.rows, model.cols, CV_8U, 255);
    return setModelByAngleCount(model, mask, angleMin, angleMax, angleCount);
}

bool RotateTemplateMatch::setModelByAngleCount(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, size_t angleCount)
{
    if(angleCount < 2 || angleCount > ANGLE_COUNT_MAX) return false;
    if(angleMax < angleMin) return false;

    m_angleMin = angleMin;
    m_angleMax = angleMax;
    m_angleCount = angleCount;
    m_angleDet = (angleMax - angleMin) / angleCount; // 最少要两个点

    return buildRotatedModels(model, mask);

}

// bool RotateTemplateMatch::setModelByDetAngle(const cv::Mat &model, double angleMin, double angleMax, double detAngle)
// {
//     cv::Mat mask = cv::Mat(model.rows, model.cols, CV_8U, 255);
//     return setModelByDetAngle(model, mask, angleMin, angleMax, detAngle);
// }

bool RotateTemplateMatch::buildRotatedModels(const cv::Mat &model, const cv::Mat &mask)
{
    if(model.empty()) return false;
    if(model.size() != mask.size()) return false;

    m_rotatedModels.clear();
    m_rotatedMasks.clear();
    m_angles.clear();
    m_scores.clear();

    m_model = model.clone();
    m_mask = mask.clone();
    m_modelSize.height = model.rows;
    m_modelSize.width = model.cols;
    for(size_t i = 0; i < m_angleCount; i++)
    {
        double angle = m_angleMin + (i + 0.5) * m_angleDet;
        m_angles.push_back(angle);
        m_scores.push_back(0.0);
        cv::Mat rotatedModel, rotatedMask;
        switch (m_sizeMode) {
        case MODEL_SIZE_MODE::FIX_MODEL_SIZE:
            imageRotateFixSize(model, angle, rotatedModel, rotatedMask);
            break;
        case MODEL_SIZE_MODE::AUTO_MODEL_SIZE:
            imageRotateAutoSize(model, angle, rotatedModel, rotatedMask);
            break;
        case MODEL_SIZE_MODE::FULL_MODEL_SIZE:
            imageRotateFullSize(model, angle, rotatedModel, rotatedMask);
            break;
        }

        if(rotatedModel.rows > m_modelSize.height) m_modelSize.height = rotatedModel.rows;
        if(rotatedModel.cols > m_modelSize.width) m_modelSize.width = rotatedModel.cols;
        m_rotatedModels.push_back(rotatedModel);
        m_rotatedMasks.push_back(rotatedMask);
    }
    return true;
}

// bool RotateTemplateMatch::setModelByDetAngle(const cv::Mat &model, const cv::Mat &mask, double angleMin, double angleMax, double detAngle)
// {
//     if(detAngle <= 0) return false;
//     if(angleMax < angleMin) return false;

//     m_angleMin = angleMin;
//     m_angleCount = ( angleMax - angleMin ) / detAngle + 0.9; //
//     m_angleDet = detAngle;
//     m_angleMax = angleMin + detAngle * m_angleCount;

//     return buildRotatedModels(model, mask);
// }

bool RotateTemplateMatch::rotatedMatch(const cv::Mat &image, cv::RotatedRect &rrt, double &score)
{
    cv::Rect rect;
    size_t index;
    bool ret = rotatedMatch(image, rect, index, score);
    if( !ret ) return false;

    double bestAngle = m_angles[index]; // 这个是原始的计算值，可以用插值做改善
    double bestX = rect.x;
    double bestY = rect.y;
//    if(m_enableSubPix)
//    {
//        angleInterpolate(index, bestAngle, score);
//    }

    double centerX = bestX + rect.width / 2.0;
    double centerY = bestY + rect.height /2.0;
    rrt = cv::RotatedRect(cv::Point2f(centerX, centerY), m_model.size(), bestAngle);
    std::cout << "RotateTemplateMatch::rotatedMatch best angle = " << bestAngle << std::endl;

    // cv::namedWindow("first match", cv::WINDOW_GUI_EXPANDED);
    // cv::Mat im2 = image.clone();
    // drawRotatedRect(im2, rrt, cv::Scalar(255, 255, 255), 2);
    // cv::imshow("first match", im2);
    return true;
}

bool RotateTemplateMatch::rotatedMatch(const cv::Mat &image, cv::Rect &rect, size_t &index, double &score)
{
    if(image.empty()) return false;
    if(m_model.empty()) return false;
    if(image.rows < m_modelSize.height || image.cols < m_modelSize.width) return false;

    double bestScore = -1e6, currentScore = -1e6;
    switch(m_matchMethod)
    {
    case cv::TM_SQDIFF:
    case cv::TM_SQDIFF_NORMED:
        bestScore = std::numeric_limits<float>::max();
         break;
    default:
        bestScore = std::numeric_limits<float>::min();
    }

    int bestIndex = 0;
    cv::Rect bestRect;

#pragma omp parallel for //num_threads(1)
    for(int index = 0; index < (int)m_angleCount; index++)
    {
        bool status = templateMatch(image, m_rotatedModels[index], m_rotatedMasks[index], m_matchMethod, rect, currentScore);
        if(status != true) continue;
        m_scores[index] = currentScore;

        #pragma omp critical
        if(better(bestScore, currentScore, m_matchMethod))
        {
            bestScore = currentScore;
            bestIndex = (int) index;
            bestRect = rect;
        }
        std::cout << "rotatedMatch: index = " << index << ", angle = " << m_angles[index] << ", score = " << currentScore << std::endl;
    }
    if(m_enableSubPix)
    {
        //TODO: 这里需要补充上亚像素插值算法
    }
    std::cout << "rotatedMatch: best score = " << bestScore << ", best angle = " << m_angles[bestIndex] << std::endl;
    rect = bestRect;
    score = bestScore; // 这个是原始的计算值，可以用插值做改善
    index = bestIndex;
    return true;
}

bool RotateTemplateMatch::singleMatch(const cv::Mat &image, int angleIndex, cv::RotatedRect &rrt, double &score)
{
    rrt = cv::RotatedRect(cv::Point2f(0, 0), cv::Size2f(0, 0), 0.0);
    score = 0.0;

    if(image.empty()) return false;
    if(m_rotatedModels[angleIndex].empty()) return false;
    if(angleIndex < 0 || angleIndex > m_angleCount) return false;
    if(image.rows < m_rotatedModels[angleIndex].rows || image.cols < m_rotatedModels[angleIndex].cols) return false;

//    std::cout << "RotateTemplateMatch::singleMatch SubPix = " << m_enableSubPix <<  ", model size = " << m_model.size() << std::endl;
    if(m_enableSubPix)
    {
        cv::Rect2f rect;
        bool status = templateMatchSubPixel(image, m_rotatedModels[angleIndex], m_rotatedMasks[angleIndex], m_matchMethod, rect, score);
        //std::cout << "status = " << status << std::endl;
        if(status != true)
        {
//            rrt = cv::RotatedRect(cv::Point2f(0, 0), cv::Size2f(0, 0), 0.0);
            return false;
        }
        double centerX = rect.x + rect.width / 2.0;
        double centerY = rect.y + rect.height /2.0;

        rrt = cv::RotatedRect(cv::Point2f(centerX, centerY), m_model.size(), m_angles[angleIndex]);
        m_scores[angleIndex] = score;
    }
    else
    {
        cv::Rect rect;
        bool status = templateMatch(image, m_rotatedModels[angleIndex], m_rotatedMasks[angleIndex], m_matchMethod, rect, score);
        if(status != true)
        {
//            rect = cv::Rect(0, 0, 0, 0);
            return false;
        }

        double centerX = rect.x + rect.width / 2.0;
        double centerY = rect.y + rect.height /2.0;
        rrt = cv::RotatedRect(cv::Point2f(centerX, centerY), m_model.size(), m_angles[angleIndex]);
        m_scores[angleIndex] = score;
    }
    return true;
}

bool RotateTemplateMatch::singleMatch(const cv::Mat &image, int angleIndex, cv::Rect &rect, double &score)
{
    rect = cv::Rect(0, 0, 0, 0);
    score = 0.0;
    if(image.empty()) return false;
    if(angleIndex < 0 || angleIndex > m_angleCount) return false;
    if(m_rotatedModels[angleIndex].empty()) return false;
    if(image.rows < m_rotatedModels[angleIndex].rows || image.cols < m_rotatedModels[angleIndex].cols) return false;

    return templateMatch(image, m_rotatedModels[angleIndex], m_rotatedMasks[angleIndex], m_matchMethod, rect, score);
}

bool RotateTemplateMatch::rectSize(size_t indexMin, size_t indexMax, int &width, int &height)
{
    width = 0;
    height = 0;
    if(indexMax < indexMin) return false;
    if(indexMax >= m_angleCount) indexMax = m_angleCount - 1;
    for(size_t index = indexMax; index <= indexMax; index ++)
    {
        if(m_rotatedModels[index].rows > height) height = m_rotatedModels[index].rows;
        if(m_rotatedModels[index].cols > width) width = m_rotatedModels[index].cols;
    }
    return true;
}

bool RotateTemplateMatch::angleInterpolate(size_t index, double &angle, double &score)
{
    if(index >= m_angleCount) return false;
    if(index == 0 || index == m_angleCount - 1)
    {
        angle = m_angles[index];
        score = m_scores[index];
        return true;
    }
    double a, b, c; //二次函数
    c = m_scores[index];
    a = - m_scores[index] + (m_scores[index - 1] + m_scores[index + 1]) / 2.0;
    b = (m_scores[index + 1] - m_scores[index - 1]) / 2.0;

//    std::cout << "a = " <<a <<", b = " <<b << ", c = " << c <<std::endl;
//    std::cout << "m_angles[index] = " << m_angles[index] << ", det = " <<  b/ (a * 2.0) * m_angleDet << std::endl;
    double amend =  - b / (a * 2.0) * m_angleDet;
    angle = m_angles[index] + amend;
    score = a * amend * amend + b * amend + c;

//    std::cout << "score[index - 1] = " << m_scores[index - 1] << ", score[index] = "  << m_scores[index] << ", score[index + 1] = " << m_scores[index + 1] << std::endl;
//    std::cout << "det angle = " << amend << ", score = " << score << std::endl;
    return true;
}

bool RotateTemplateMatch::rotatedModel(size_t angleIndex, cv::Mat &model)
{
    if(angleIndex > m_angleCount - 1) return false;
    model = m_rotatedModels[angleIndex];
    return true;
}

bool RotateTemplateMatch::rotatedMask(size_t angleIndex, cv::Mat &mask)
{
    if(angleIndex > m_angleCount - 1) return false;
    mask = m_rotatedMasks[angleIndex];
    return true;
}

}
